package cn.feizhu.aimarket.ai.workflow.node;

import cn.feizhu.aimarket.ai.workflow.model.constant.WorkflowControlConstant;
import cn.feizhu.aimarket.ai.workflow.model.enums.NodeStatus;
import com.alibaba.cloud.ai.graph.NodeOutput;
import com.alibaba.cloud.ai.graph.OverAllState;
import com.alibaba.cloud.ai.graph.action.NodeAction;
import com.alibaba.cloud.ai.graph.async.AsyncGenerator;
import com.alibaba.cloud.ai.graph.streaming.StreamingChatGenerator;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatResponse;
import reactor.core.publisher.Flux;

import java.util.Map;
import java.util.Objects;

/**
 * 大纲生成节点
 * 基于处理后的指令生成详细的内容大纲
 *
 * @author feizhu
 * @since 2025/10/16
 */
public class OutlineGenerateNode implements NodeAction {

    private static final Logger logger = LoggerFactory.getLogger(OutlineGenerateNode.class);

    private final ChatClient outLineChatClient;

    private final Map<String, NodeStatus> node2Status;

    public static final String NODE_NAME = "outline_generate";

    public OutlineGenerateNode(ChatClient outLineChatClient, Map<String, NodeStatus> node2Status) {
        this.outLineChatClient = outLineChatClient;
        this.node2Status = node2Status;
    }

    @Override
    public Map<String, Object> apply(OverAllState state) {
        node2Status.put(NODE_NAME, NodeStatus.RUNNING);
        logger.info("=== 大纲生成节点开始执行 ===");

        try {
            // 获取处理后的指令
            String processedCommand = (String) state.value(WorkflowControlConstant.PROCESSED_COMMAND)
                    .orElse("未找到处理后的指令信息");

            logger.info("基于处理后的指令生成大纲: {}", processedCommand);

            // 创建流式响应
            Flux<ChatResponse> chatResponseFlux = this.outLineChatClient.prompt()
                    .user(processedCommand)
                    .stream()
                    .chatResponse();

            // 使用 StreamingChatGenerator 构建异步生成器
            AsyncGenerator<? extends NodeOutput> generator = StreamingChatGenerator.builder()
                    .startingNode(NODE_NAME) // 这个不加就会导致流式的时候没有名字
                    .startingState(state)
                    .mapResult(response -> {
                        String text = response.getResult().getOutput().getText();
//                        logger.info("生成的大纲: {}", text);
                        node2Status.put(NODE_NAME, NodeStatus.COMPLETED);
                        return Map.of(WorkflowControlConstant.OUTLINE, text);
                    })
                    .build(chatResponseFlux);

            return Map.of(WorkflowControlConstant.OUTLINE, generator);

        } catch (Exception e) {
            logger.error("大纲生成节点执行异常", e);
            node2Status.put(NODE_NAME, NodeStatus.FAILED);
            return Map.of(WorkflowControlConstant.OUTLINE, "大纲生成失败：" + e.getMessage());
        } finally {
            logger.info("=== 大纲生成节点执行完成 ===");
        }
    }
}
